#! /usr/bin/env python
# -*- coding: utf-8 -*-

# *************************************************************
#     Filename @  retry.py
#       Author @  Huoty
#  Create date @  2016-08-04 19:58:39
#  Description @  Retry again if an exception occurs
# *************************************************************

"""异常重试装饰器"""

import time
import random
import traceback
import functools


class Logger(object):
    """日志处理器

    初始化 logger 为 None 时则不输出日志，仅支持 error、warn、info 三个日志级别
    """
    def __init__(self, logger=None):
        self.logger = logger

    def __log_swith(self, log_type, msg):
        {
            "error": getattr(self.logger, "error"),
            "warn": getattr(self.logger, "warning") if hasattr(self.logger, "warning") else getattr(self.logger, "warn"),
            "info": getattr(self.logger, "info"),
        }[log_type](msg) if self.logger else None

    def error(self, msg):
        self.__log_swith("error", msg)

    def warn(self, msg):
        self.__log_swith("warn", msg)

    def info(self, msg):
        self.__log_swith("info", msg)


class Retrier(object):
    """重试器

    初始化参数说明：

        tries: 重试次数, 默认一直重试
        delay: 两次重试之间的间隔时间, 可以为一个 list 或 tuple 来指定每次间隔的时间, 且重试次数为其长度
        exceptions: 指定哪些异常发生时需要重试, 可以是一个 list 或者 tuple
        exponential, max_delay: 时间间隔按指数增长, 如果 exponential 为 True, 则每次的时间间隔为：
            (2 ** previous_tried_count) * delay
        random_delay_min, random_delay_max: 随机间隔时间
        delay_func: 间隔函数, 每次重试会调用该函数产生一个间隔时间, 调用该函数时会把第几次重试作为参数传入
        logger: 日志系统, 为 None 时不输出信息

    间隔时间的类型优先级为: 函数间隔 > 列表间隔 > 指数间隔 > 随机间隔 > 固定间隔
    即如果指定了 exponential 为 True, 又同时指定了 random_delay_min, random_delay_max, 则默认采用指数间隔
    """

    def __init__(self,
                 tries=None,
                 delay=None,
                 exceptions=Exception,
                 exponential=False, max_delay=None,
                 random_delay_min=None, random_delay_max=None,
                 delay_func=None,
                 logger=None):

        if exponential and not isinstance(delay, (int, float)):
            raise TypeError("当指定 exponential 为 True 时, delay 必须为一个 int 或者 float 类型")

        self.tries = int(tries) if tries else -1
        self.delay = (delay if isinstance(delay, (list, tuple)) else float(delay)) if delay else 0
        self.exceptions = tuple(exceptions) if isinstance(exceptions, list) else exceptions
        self.exponential = exponential
        self.max_delay = float(max_delay) if max_delay is not None else float("inf")
        self.random_delay_min = random_delay_min
        self.random_delay_max = random_delay_max
        self.delay_func = delay_func
        self.log = Logger(logger)

    @staticmethod
    def __add_suffix(num):
        num = str(num)
        suffix = {"1": "st", "2": "nd", "3": "rd"}
        suffix.update({n: "th" for n in list("0456789")})
        return num + suffix[num[-1]]

    def calc_delay(self, tried_times):
        if isinstance(self.delay, (list, tuple)):
            return float(self.delay[tried_times])
        elif self.exponential:
            delay = self.delay * (2 ** tried_times)
            return delay if delay < self.max_delay else self.max_delay
        elif self.random_delay_min or self.random_delay_max:
            random_delay_min = 0 if self.random_delay_min is None else float(self.random_delay_min)
            random_delay_max = 1 << 64 if self.random_delay_max is None else float(self.random_delay_min)
            return random.uniform(random_delay_min, random_delay_max)
        else:
            return self.delay

    def call(self, func, *args, **kwargs):
        tries = len(self.delay) if isinstance(self.delay, (list, tuple)) else self.tries
        exceptions = self.exceptions
        log = self.log
        try_count = -1
        while 1:
            try:
                try_count += 1
                return func(*args, **kwargs)
            except exceptions as e:
                if tries == 0:
                    raise

                tries -= 1
                log.error("%s(%s, %s) failed, try_count=%s, exception=%s" % \
                    (func.__name__, args, kwargs, try_count, traceback.format_exc()))
                delay = self.delay_func(try_count+1) if self.delay_func else self.calc_delay(try_count)
                time.sleep(delay)
                log.info("due to exception, will be carried out %s try" % self.__add_suffix(try_count+1))
        pass


def retry(*dargs, **dkw):
    if len(dargs) == 1 and callable(dargs[0]):

        def wrapper_simple(func):
            @functools.wraps(func)
            def wrapped_func(*args, **kw):
                return Retrier().call(func, *args, **kw)
            return wrapped_func
        return wrapper_simple(dargs[0])

    else:

        def wrapper(func):
            @functools.wraps(func)
            def wrapped_func(*args, **kw):
                return Retrier(*dargs, **dkw).call(func, *args, **kw)
            return wrapped_func
        return wrapper


# Script starts from here

if __name__ == "__main__":
    pass
